'''
Created on Sep 18, 2009

@author: mkiyer
'''
from veggie.sample.sampledb2 import get_sampledb
from veggie.sample.samplegroup import SampleGroup

import collections
import logging
import numpy as np

def get_datasets_from_sample_group(sgroup, covdb):
    '''
    return a dictionary where key is sample name and value is 
    a list of hdf5 coverage datasets
    
    this function actually gets the hdf5 datasets, whereas the previous
    version simply returns the dataset names
    '''
    dsets = collections.defaultdict(lambda: [])
    for s in sgroup.samples:
        sample_dsets = get_datasets_from_sample_name(s, covdb)
        if len(sample_dsets) > 0:
            logging.debug('    added sample: %s' % s)
            dsets[s] = sample_dsets
        else:
            logging.warning('    %s has no coverage lanes.. skipping' % s)
    return dsets

def get_datasets_from_sample_name(sample_name, covdb):
    dsets = []
    for fc_lane in get_sampledb().params[sample_name]['lanes']:
        fc_lane_id = '_'.join(map(str, fc_lane))
        if fc_lane_id in covdb:
            logging.debug('    added: %s lane: %s' % (sample_name, fc_lane_id))
            dsets.append(covdb[fc_lane_id])
        else:
            logging.warning('    sample: %s lane: %s not found... skipping' % (sample_name, fc_lane_id))
    return dsets

class SampleGroupCoverage(object):
    '''
    this pairs a SampleGroup object with coverage data stored in a separate location
    '''    
    def __init__(self, sample_group, coverage_dsets):
        self.dsets = get_datasets_from_sample_group(sample_group, coverage_dsets)
        self._sample_group = sample_group
    def __len__(self):
        return len(self.dsets)
    @property
    def name(self):
        return self._sample_group.name
    @property
    def color(self):
        return self._sample_group.color
    @property
    def samples(self):
        return self.dsets.keys()
    def itersamples(self):
        return self.dsets.iterkeys()
    def iterkeys(self):
        return self.dsets.iterkeys()
    def iteritems(self):
        return self.dsets.iteritems()
    def itervalues(self):
        return self.dsets.itervalues()
    
    def calc_avg_cov_rpkm(self, chrom, start, end):
        for sample_name, dset_list in self.dsets.iteritems():
            yield sample_name, calc_avg_cov_rpkm(chrom, start, end, dset_list)
    
    

def calc_avg_cov_rpkm(chrom, start, end, dsets):
    '''
    calculate the average coverage from the datasets specified 
    over the interval specified by (chrom, start, end)
    
    returns the average coverage of the interval in RPKM
    '''
    avgcov = 0.0
    for dset in dsets:
        avgcov += calc_rpkm_array(dset[chrom][start:end])
    return avgcov / len(dsets)


def calc_avg_cov_array(chrom, start, end, dsets):
    '''
    calculate the average coverage from the datasets specified 
    over the interval specified by (chrom, start, end)
    
    returns an array over the entire interval
    '''
    covarray = np.zeros(end-start, dtype=np.float)
    for dset in dsets:
        covarray += dset[chrom][start:end]
    covarray /= len(dsets)
    return covarray

def calc_avg_cov_arrays(chrom, start, end, sample_groups):
    '''
    call calc_avg_cov_array for every sample in the sample_groups
    and return a dictionary where the key is the sample name and
    the value is the coverage
    ''' 
    sample_cov_arrays = {}     
    for sgroup in sample_groups:
        for sample_name, sample_dsets in sgroup.iteritems():
            avgcovarr = calc_avg_cov_array(chrom, start, end, sample_dsets)                
            sample_cov_arrays[sample_name] = avgcovarr
    return sample_cov_arrays

def calc_rpkm(chrom, start, end, dset):
    return calc_rpkm_array(dset[chrom][start:end])

def calc_rpkm_array(covarray):
    return 1000.0 * np.sum(covarray) / covarray.size

def calc_rpkm_matrix(covmat, axis=0):
    return (1000.0 / covmat.shape[axis]) * np.sum(covmat, axis=axis)

#def gen_avg_cov_array(start, end, chrom_dsets):
#    # get average coverage for this interval
#    covarrays = np.zeros((end - start, len(chrom_dsets)), dtype=np.float)
#    for j, dsets in enumerate(chrom_dsets.itervalues()):
#        # get average coverage for all lanes of this sample
#        for dset in dsets:
#            covarrays[:,j] += dset[start:end]
#        covarrays[:,j] /= len(dsets)
#    return covarrays
